#ifndef HIR_DFS_HPP
#define HIR_DFS_HPP
#pragma once

namespace hir
{

class BasicBlockFunctor
{
public:
    typedef BasicBlock::BlockList::iterator iterator;
    inline BasicBlock * iteratorToObject(iterator & i)
    {
        return *i;
    }
};

class DominateFunctor : public BasicBlockFunctor
{
public:
    BasicBlock::BlockList & getBlockList(BasicBlock * bb)
    {
        return bb->dominatee();
    }
};

class SuccessorFunctor : public BasicBlockFunctor
{
public:
    BasicBlock::BlockList & getBlockList(BasicBlock * bb)
    {
        return bb->successor();
    }
};

class PredecessorFunctor : public BasicBlockFunctor
{
public:
    BasicBlock::BlockList & getBlockList(BasicBlock * bb)
    {
        return bb->predecessor();
    }
};

template <typename DataType,typename Functor>
class DFSPreorderIterator
{
public:
    typedef DataType datatype;
    typedef Functor functor;
    DFSPreorderIterator() : bb_(nullptr)
        , functor_(functor())
        , data_(datatype())
    {
    };

    explicit DFSPreorderIterator(BasicBlock * bb,datatype && data = datatype(),functor && f = functor()): bb_(bb)
        , functor_(f)
        , data_(data)
    {
        if(bb_)
            bitmap_.set(bb->id());
    }

    bool operator == (const DFSPreorderIterator & other)
    {
        return other.bb_ == this->bb_;
    }

    bool operator != (const DFSPreorderIterator & other)
    {
        return other.bb_ != this->bb_;
    }

    inline BasicBlock * operator *()
    {
        return bb_;
    }

    datatype & data()
    {
        return data_;
    }

    DFSPreorderIterator & operator ++ ()
    {
        if(!bb_)
        {
            // it's the end
            return *this;
        }

        if(functor_.getBlockList(bb_).size() == 0)
        {
            // has nothing!
            return currentBbExhaustedCase();
        }

        // bb_ still has something
        typename functor::iterator begin = functor_.getBlockList(bb_).begin();
        BasicBlock * willBeBeg = functor_.iteratorToObject(begin);
        while(begin != functor_.getBlockList(bb_).end() && bitmap_.test(willBeBeg->id()))
        {
            // this node has been visited
            willBeBeg = functor_.iteratorToObject(begin);
            ++begin;
        }
        if(begin == functor_.getBlockList(bb_).end())
        {
            // exhausted go to exhuasted case
            return currentBbExhaustedCase();
        }

        bitmap_.set(willBeBeg->id());
        if(begin == functor_.getBlockList(bb_).end())
        {
            // will not push state in this case
            // just one to iterate
            bb_ = willBeBeg;
            return *this;
        }
        State s = {bb_ , data_, begin };
        stateStack_.push_back(s);
        bb_ = willBeBeg;
        // data_ = data_;
        return *this;
    }
private:

    DFSPreorderIterator & currentBbExhaustedCase()
    {
        if(stateStack_.size() == 0)
        {
            // it's the end!
            bb_ = nullptr;
            return *this;
        }
        // forward or pop
        while(stateStack_.size() > 0)
        {
            State & state = stateStack_.back();
            BasicBlock * upper = state.bb_;
            // check end of the state
            if(functor_.getBlockList(bb_).end() == state.workingIterator_)
            {
                // end of the state
                // pop it
                stateStack_.pop_back();
                continue;
            }
            // try forward
            while(functor_.getBlockList(upper).end() != state.workingIterator_)
            {
                BasicBlock * willBeBeg = functor_.iteratorToObject(state.workingIterator_);
                if(bitmap_.test(willBeBeg->id()))
                {
                    // has been visited
                    ++state.workingIterator_;
                    continue;
                }
                bitmap_.set(willBeBeg->id());
                ++state.workingIterator_;
                bb_ = willBeBeg;
                // inherit parent's data
                data_ = state.data_;
                // now end this
                return *this;
            }
            // check if state is end
            if(functor_.getBlockList(upper).end() == state.workingIterator_)
            {
                // end of the state
                // pop it
                stateStack_.pop_back();
                continue;
            }
            // forward success
            return *this;
        }
        DEBUG_ASSERT(stateStack_.size() == 0);
        // stateStack dry means end
        bb_ = nullptr;
        return *this;
    }

    struct State
    {
        BasicBlock * bb_;
        datatype data_;
        typename functor::iterator workingIterator_;
    };
    BasicBlock * bb_;
    datatype data_;
    typedef std::list<State> StateStack;
    StateStack stateStack_;
    std::bitset<MAX_BASIC_BLOCK_ID> bitmap_;
    functor functor_;
};

typedef DFSPreorderIterator<int,SuccessorFunctor> SuccessorDFSPreorderIterator;
}

#endif /* HIR_DFS_HPP */
